"""
MCP (Model Context Protocol) router implementation
Handles JSON-RPC 2.0 requests for MCP protocol
"""
import json
import time
import uuid
from typing import Dict, Any, List, Optional
from fastapi import APIRouter, HTTPException, Request, Response
from fastapi.responses import JSONResponse
import asyncio
import logging

from ..models.mcp import (
    MCPRequest, MCPResponse, MCPError, MCPMethod,
    ToolDefinition, ToolsListResult, ToolCallParams, ToolCallResult,
    InitializeParams, InitializeResult, ServerCapabilities, MCPNotification
)

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/mcp", tags=["MCP Protocol"])


class ToolRegistry:
    """Registry for available tools"""
    
    def __init__(self):
        self._tools: Dict[str, Dict[str, Any]] = {}
        self._register_default_tools()
    
    def _register_default_tools(self):
        """Register default available tools"""
        self.register_tool(
            name="echo",
            description="Echo back the input text",
            input_schema={
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "Text to echo back"
                    }
                },
                "required": ["text"]
            },
            handler=self._echo_tool
        )
        
        self.register_tool(
            name="get_time",
            description="Get current timestamp",
            input_schema={
                "type": "object",
                "properties": {},
                "required": []
            },
            handler=self._get_time_tool
        )
    
    def register_tool(self, name: str, description: str, input_schema: Dict[str, Any], handler):
        """Register a tool with its handler"""
        self._tools[name] = {
            "description": description,
            "input_schema": input_schema,
            "handler": handler
        }
    
    def get_tools_list(self) -> List[ToolDefinition]:
        """Get list of available tools"""
        return [
            ToolDefinition(
                name=name,
                description=tool["description"],
                input_schema=tool["input_schema"]
            )
            for name, tool in self._tools.items()
        ]
    
    async def call_tool(self, name: str, arguments: Dict[str, Any]) -> ToolCallResult:
        """Call a tool by name with arguments"""
        if name not in self._tools:
            return ToolCallResult(
                content=[{"type": "text", "text": f"Tool '{name}' not found"}],
                isError=True
            )
        
        try:
            handler = self._tools[name]["handler"]
            if asyncio.iscoroutinefunction(handler):
                result = await handler(arguments)
            else:
                result = handler(arguments)
            
            return ToolCallResult(
                content=[{"type": "text", "text": str(result)}],
                isError=False
            )
        except Exception as e:
            logger.error(f"Error calling tool '{name}': {str(e)}")
            return ToolCallResult(
                content=[{"type": "text", "text": f"Error: {str(e)}"}],
                isError=True
            )
    
    async def _echo_tool(self, arguments: Dict[str, Any]) -> str:
        """Echo tool implementation"""
        return arguments.get("text", "")
    
    async def _get_time_tool(self, arguments: Dict[str, Any]) -> str:
        """Get current time tool implementation"""
        return f"Current timestamp: {int(time.time())}"


# Global tool registry
tool_registry = ToolRegistry()


def create_error_response(request_id: Optional[str], code: int, message: str, data: Any = None) -> MCPResponse:
    """Create an error response"""
    return MCPResponse(
        id=request_id,
        error=MCPError(code=code, message=message, data=data)
    )


def create_success_response(request_id: Optional[str], result: Any) -> MCPResponse:
    """Create a success response"""
    return MCPResponse(id=request_id, result=result)


@router.post("/", response_model=MCPResponse)
async def handle_mcp_request(request_data: Dict[str, Any]) -> Response:
    """
    Main MCP endpoint handler
    Processes JSON-RPC 2.0 requests according to MCP protocol
    """
    try:
        # Parse the request
        try:
            mcp_request = MCPRequest(**request_data)
        except Exception as e:
            logger.error(f"Invalid request format: {str(e)}")
            error_response = create_error_response(
                request_id=request_data.get("id"),
                code=-32600,  # Invalid Request
                message="Invalid Request",
                data=str(e)
            )
            return JSONResponse(content=error_response.model_dump())
        
        # Handle different MCP methods
        if mcp_request.method == MCPMethod.INITIALIZE:
            response = await handle_initialize(mcp_request)
        elif mcp_request.method == MCPMethod.TOOLS_LIST:
            response = await handle_tools_list(mcp_request)
        elif mcp_request.method == MCPMethod.TOOLS_CALL:
            response = await handle_tools_call(mcp_request)
        else:
            logger.warning(f"Unknown method: {mcp_request.method}")
            response = create_error_response(
                request_id=mcp_request.id,
                code=-32601,  # Method not found
                message=f"Method '{mcp_request.method}' not found"
            )
        
        return JSONResponse(content=response.model_dump())
    
    except Exception as e:
        logger.error(f"Internal error processing MCP request: {str(e)}")
        error_response = create_error_response(
            request_id=request_data.get("id"),
            code=-32603,  # Internal error
            message="Internal error",
            data=str(e)
        )
        return JSONResponse(content=error_response.model_dump())


async def handle_initialize(request: MCPRequest) -> MCPResponse:
    """Handle initialize method"""
    try:
        if request.params:
            params = InitializeParams(**request.params)
            logger.info(f"Client initialized with protocol version: {params.protocolVersion}")
        
        result = InitializeResult(
            protocolVersion="2024-11-05",
            capabilities=ServerCapabilities(
                tools={"listChanged": True}
            ),
            serverInfo={
                "name": "TeachWorkflow MCP Server",
                "version": "1.0.0"
            }
        )
        
        return create_success_response(request.id, result.model_dump())
    
    except Exception as e:
        logger.error(f"Error in initialize: {str(e)}")
        return create_error_response(
            request_id=request.id,
            code=-32603,
            message="Initialize failed",
            data=str(e)
        )


async def handle_tools_list(request: MCPRequest) -> MCPResponse:
    """Handle tools/list method"""
    try:
        tools = tool_registry.get_tools_list()
        result = ToolsListResult(tools=tools)
        
        logger.info(f"Returning {len(tools)} available tools")
        return create_success_response(request.id, result.model_dump())
    
    except Exception as e:
        logger.error(f"Error in tools/list: {str(e)}")
        return create_error_response(
            request_id=request.id,
            code=-32603,
            message="Failed to list tools",
            data=str(e)
        )


async def handle_tools_call(request: MCPRequest) -> MCPResponse:
    """Handle tools/call method"""
    try:
        if not request.params:
            return create_error_response(
                request_id=request.id,
                code=-32602,  # Invalid params
                message="Missing parameters for tools/call"
            )
        
        params = ToolCallParams(**request.params)
        logger.info(f"Calling tool: {params.name} with arguments: {params.arguments}")
        
        result = await tool_registry.call_tool(params.name, params.arguments)
        
        return create_success_response(request.id, result.model_dump())
    
    except Exception as e:
        logger.error(f"Error in tools/call: {str(e)}")
        return create_error_response(
            request_id=request.id,
            code=-32603,
            message="Tool call failed",
            data=str(e)
        )


@router.get("/tools", response_model=ToolsListResult)
async def get_available_tools():
    """
    Get list of available tools (HTTP GET endpoint for convenience)
    """
    tools = tool_registry.get_tools_list()
    return ToolsListResult(tools=tools)


@router.post("/tools/{tool_name}/call")
async def call_tool_direct(tool_name: str, arguments: Dict[str, Any]):
    """
    Direct tool call endpoint (HTTP POST for convenience)
    """
    result = await tool_registry.call_tool(tool_name, arguments)
    if result.isError:
        raise HTTPException(status_code=400, detail=result.content[0].get("text", "Tool call failed"))
    return {"result": result.content}


# Export the tool registry for use by other modules
__all__ = ["router", "tool_registry", "ToolRegistry"]